import numpy

import charmicat
from charmicat import Cut

def all4g_pos_str(c, format_str):
    columns = 'pi01_g1{0} pi01_g2{0} pi02_g1{0} pi02_g2{0}'.format(c).split()
    return ' & '.join([format_str.format(c) for c in columns])

e9oe25_cuts = [charmicat.Cut(name='pi01_g1e9oe25',
                             pos_str='pi01_g1e9oe25 > 0.9'),
               charmicat.Cut(name='pi01_g2e9oe25',
                             pos_str='pi01_g2e9oe25 > 0.9'),               
               charmicat.Cut(name='pi02_g1e9oe25',
                             pos_str='pi02_g1e9oe25 > 0.9'),
               charmicat.Cut(name='pi02_g2e9oe25',
                             pos_str='pi02_g2e9oe25 > 0.9'),
               ]


all4g_e9oe25_gt_095 = charmicat.cut_intersection(e9oe25_cuts,
                                                 name='all4g_e9oe25_gt_0.95')

reduction = [charmicat.Cut(name='pi0pi0_chisq_lt_20',
                           pos_str="((pi01_chisq + pi02_chisq) < 20)",
                           ),
             charmicat.Cut(name='all4g_energy_gt_0.05',
                           pos_str=all4g_pos_str('e', '({0} > 0.05)'),
                           ),
             charmicat.Cut(name='delta_m',
                           pos_str='(abs(mm_off_pi0pi0 - m_mumu) < 0.2)',
                           ),
             charmicat.Cut(name='p_cm_total_lt_0.250',
                          pos_str='(p_cm_total < 0.250)',
                          ),
             ]

krokovny = [charmicat.Cut(name='all4g_energy_gt_0.05',
                          pos_str=all4g_pos_str('e', '({0} > 0.05)'),
                          ),
            charmicat.Cut(name='missing_e_lt_0.200',
                          pos_str='(abs(missing_e) < 0.2)',
                          ),
            charmicat.Cut(name='p_cm_total_lt_0.250',
                          pos_str='(p_cm_total < 0.250)',
                          ),
            charmicat.Cut(name='m_mumu_gt_0.150',
                          pos_str='(m_mumu > 0.150)',
                          ),
            ]

misc_cuts = [charmicat.Cut(name='m_mumu_y1s_2sig',
                           pos_str='(m_mumu > 9.3) & (m_mumu < 9.59)',
                           ),
             charmicat.Cut(name='m_mumu_y2s',
                           pos_str='(m_mumu > 9.8) & (m_mumu < 10.2)',
                           ),
             charmicat.Cut(name='m_mumu_y3s',
                           pos_str='(m_mumu > 10.2) & (m_mumu < 10.6)',
                           ),
             charmicat.Cut(name='m_mumu_bt_1and2.5',
                           pos_str='(m_mumu > 1.0) & (m_mumu < 2.5)',
                           ),
             charmicat.Cut(name='fwr2_lt_0.35',
                           pos_str='r2 < 0.35',
                           ),        
             charmicat.Cut(name='muon_llh_gt_0.1',
                           pos_str='((mup_muid_llh > 0.1) & (mum_muid_llh > 0.1))',
                           desc='Events where both charged tracks have muid > 0.1',
                           ),
             charmicat.Cut(name='delta_m',
                           pos_str='(abs(m_mumu - mm_off_pi0pi0) < 0.1)',
                           desc='Events with |m_mumu - mm_off_pi0pi0| < 100 MeV'
                           ),     
             charmicat.Cut(name='one_plus_one_minus',
                           pos_str='(n_pm_combos == 1)',
                           desc='Events with one positive track and one negative track',
                           ),
             charmicat.Cut(name='exclusive_to_15MeV',
                           pos_str='(abs(2*beam_e - m_mumupi0pi0) < 0.015)',
                           desc='Events where |Ecm - m_mumupi0pi0| < 15 MeV',
                           ),
             charmicat.Cut(name='delta_m_lt_0.150',
                           pos_str='(abs(m_mumu - mm_off_pi0pi0) < 0.150)',
                           desc='Events where |m_mumu - mm_off_pi0pi0| < 150 MeV',
                           ),
             charmicat.Cut(name='hbbgm_delta_m_lt_0.150',
                           pos_str='(abs(m2mu - recm2pi) < 0.150)',
                           desc='Events where |m_mumu - mm_off_pi0pi0| < 150 MeV',
                           ),
             charmicat.Cut(name='hbbgm_m_mumu_y123s',
                           pos_str='((abs(m2mu - 9.460) < 0.150) | (abs(m2mu - 10.023) < 0.150) | (abs(m2mu - 10.355) < 0.150))',
                           desc='Events with m_mumu within 150 MeV of the nominal Y(nS) (n=1,2,3) mass',
                           ),
             charmicat.Cut(name='y3s_2d',
                           pos_str="(abs(m2mu - 10.355) < 0.150) & (abs(recm2pi - 10.355) < 0.150)",
                           desc=''),
             charmicat.Cut(name='bgm_m_mumu_y123s',
                           pos_str='((abs(m_mumu - 9.460) < 0.150) | (abs(m_mumu - 10.023) < 0.150) | (abs(m_mumu - 10.355) < 0.150))',
                           desc='Events with m_mumu within 150 MeV of the nominal Y(nS) (n=1,2,3) mass',
                           ),             
             ]



standard_cuts = [charmicat.Cut(name='pi0_mass_within_3sig',
                               pos_str="(abs(pi01_m - 0.135) < 0.0054*3) & (abs(pi02_m - 0.135) < 0.0054*3)",
                               desc='Events where both pi0 candidates have mass on (120, 150) MeV',
                               ),
                 charmicat.Cut(name='all4g_energy_gt_0.03',
                               pos_str=all4g_pos_str('e', '({0} > 0.03)'),
                               desc='Events where all four gammas are greater than 30 MeV in energy',
                               ),
                 charmicat.Cut(name='muon_mu_to_e_gt_0.6',
                               pos_str='((mup_muon_vs_electron_llh > 0.6) & (mum_muon_vs_electron_llh > 0.6))',
                               desc='Events where the two charged tracks both have mu vs e llhood > 0.6',
                               ),
                 charmicat.Cut(name='all4g_e9oe25_gt_0.85',
                               pos_str=all4g_pos_str('e9oe25', '({0} > 0.85)'),
                               desc='Events where all 4 gammas have e9/e25 > 0.85',
                               ),
                 charmicat.Cut(name='all4g_theta_lab_bt_17and150deg',
                               pos_str=all4g_pos_str('theta_lab', '(({0} > 17) & ({0} < 150))'),
                               desc='Events where all 4 gammas have theta on (17, 150) deg',
                               ),
                 charmicat.Cut(name='kill351291',
                               pos_str=('(expno != 71) | (runno != 329) | (eventno !=351291)'),
                               desc='Events where all 4 gammas have theta on (17, 150) deg',
                               ),
                 ]

bgm_cuts = [charmicat.Cut(name='muon_llh_gt_0.9',
                          pos_str='((mup_muid_llh > 0.9) & (mum_muid_llh > 0.9))',
                          desc='Events where both charged tracks have muid > 0.9',
                          ),
            charmicat.Cut(name='all4g_energy_gt_0.1',
                          pos_str=all4g_pos_str('e', '({0} > 0.1)'),
                          desc='Events where all four gammas are greater than 100 MeV in energy',
                          ),
            charmicat.Cut(name='pi0_mass_within_3sig',
                          pos_str="(abs(pi01_m - 0.135) < 0.0054*3) & (abs(pi02_m - 0.135) < 0.0054*3)",
                          desc='Events where both pi0 candidates have mass on (118, 151.2) MeV',
                          ),
            charmicat.Cut(name='all4g_e9oe25_gt_0.95',
                          pos_str=all4g_pos_str('e9oe25', '({0} > 0.95)'),
                          desc='Events where all 4 gammas have e9/e25 > 0.95',
                          ),
            charmicat.Cut(name='m_pi0pi0_gt_350MeV',
                          pos_str='(m_pi0pi0 > 0.35)',
                          desc='Events where M(pi0pi0) exceeds 350 MeV',
                          ),
            charmicat.Cut(name='mm_off_mumu_gt_300MeV',
                          pos_str='(mm_off_mumu > 0.3)',
                          desc='Events where MM(mu+mu-) exceeds 300 MeV',
                          ),
             charmicat.Cut(name='mbc_exclusive_to_15MeV',
                           pos_str='(abs(2*beam_e - mbc_mumupi0pi0) < 0.015)',
                           desc='Events where |Ecm - M_bc(mumupi0pi0)| < 15 MeV',
                           ),
            ]


hbbgm_cuts = [charmicat.Cut(name='ntr_eq_2',
                            pos_str='ntr == 2',
                            desc=''),
              charmicat.Cut(name='m2pi2mu_gt_0',
                            pos_str='m2pi2mu > 0',
                            desc=''),
              charmicat.Cut(name='delta_m',
                            pos_str='abs(sqrt(m2pi2mu) - 2*energy) < 0.015',
                            desc=''),
              charmicat.Cut(name='m2pi_gt_0.35',
                            pos_str='m2pi > 0.35',
                            desc=''),
              charmicat.Cut(name='recm2mu_gt_0.3',
                            pos_str='recm2mu > 0.3',
                            desc=''),
              ]
              
              
              

e_gamma_cuts = {}

window_250MeV = zip(numpy.arange(0.05, 0.66, 0.05), numpy.arange(0.3, 0.91, 0.05))

pi0s = 'pi01_g1e pi01_g2e pi02_g1e pi02_g2e'.split()
for lo,hi in window_250MeV:
    for gamma in pi0s:
        name = '{0}_bt_{1:03d}and{2:03d}MeV'.format(gamma, int(lo*1000), int(hi*1000))
        pos_str = '(({0} > {1}) & ({0} < {2}))'.format(gamma, lo, hi)        

        e_gamma_cuts[name] = charmicat.Cut(name, pos_str)

    name = 'all4g_bt_{1:03d}and{2:03d}MeV'.format(gamma, int(lo*1000), int(hi*1000))
    pos_str = ' & '.join(['(({0} > {1}) & ({0} < {2}))'.format(gamma, lo, hi) for gamma in pi0s])

    e_gamma_cuts[name] = charmicat.Cut(name, pos_str)

e_gamma = [charmicat.Cut(name='pi01_g1e_gt_0.125',
                         pos_str='pi01_g1e > 0.125',
                         ),
           charmicat.Cut(name='pi01_g2e_gt_0.1',
                         pos_str='pi01_g2e > 0.1',
                         ),
           charmicat.Cut(name='pi02_g1e_gt_0.75',
                         pos_str='pi02_g1e > 0.75',
                         ),
           charmicat.Cut(name='pi02_g2e_gt_0.5',
                         pos_str='pi02_g2e > 0.5',
                         ),
           ]

e_gamma_windows = [charmicat.Cut(name='pi01_g1e_bt_0.075and0.4',
                                 pos_str='(pi01_g1e > 0.075) & (pi01_g1e < 0.4)',
                                 ),
                   charmicat.Cut(name='pi01_g2e_0.025and0.2',
                                 pos_str='(pi01_g2e > 0.025) & (pi01_g2e < 0.2)',
                                 ),
                   charmicat.Cut(name='pi02_g1e_bt_0.075and0.4',
                                 pos_str='(pi02_g1e > 0.075) & (pi02_g1e < 0.4)',
                                 ),
                   charmicat.Cut(name='pi02_g2e_0.025and0.2',
                                 pos_str='(pi02_g2e > 0.025) & (pi02_g2e < 0.2)',
                                 ),
                   ]



cuts = dict([(c.name, c) for c in reduction])
#cuts.update(e_gamma_cuts)
cuts.update(dict([(c.name, c) for c in krokovny]))
cuts.update(dict([(c.name, c) for c in misc_cuts]))
cuts.update(dict([(c.name, c) for c in e_gamma]))
cuts.update(dict([(c.name, c) for c in e9oe25_cuts]))
cuts.update(dict([(c.name, c) for c in e_gamma_windows]))
cuts.update(dict([(c.name, c) for c in standard_cuts]))
cuts.update(dict([(c.name, c) for c in bgm_cuts]))
cuts.update(dict([(c.name, c) for c in hbbgm_cuts]))

cutsets = {'e_gamma':           [c.name for c in e_gamma],
           'e_gamma_windows':   [c.name for c in e_gamma_windows],
           'reduction':         [c.name for c in reduction],
           'krokovny':          [c.name for c in krokovny],
           'e9oe25':            [c.name for c in e9oe25_cuts],
           'standard':          [c.name for c in standard_cuts],
           'standard_no_muon':  [c.name for c in standard_cuts if 'muon' not in c.name],
           'bgm':               [c.name for c in bgm_cuts],
           'bgm_hbcb':          [c.name for c in bgm_cuts[-3:]],
           'hbbgm':             [c.name for c in hbbgm_cuts],
           'hbbgm_red':         [c.name for c in hbbgm_cuts] + ['hbbgm_m2mu_y123s'],
           'hbbgm_blue':        [c.name for c in hbbgm_cuts] + ['hbbgm_delta_m_lt_0.150']           
           }

cutsets['__default__'] = 'standard'



